{
  "cells": [
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "7331ed20-4a7e-4ee3-8454-528ac48315f0",
      "metadata": {
        "id": "7331ed20-4a7e-4ee3-8454-528ac48315f0"
      },
      "outputs": [],
      "source": [
        "# Copyright 2024 Google LLC\n",
        "#\n",
        "# Licensed under the Apache License, Version 2.0 (the \"License\");\n",
        "# you may not use this file except in compliance with the License.\n",
        "# You may obtain a copy of the License at\n",
        "#\n",
        "#     https://www.apache.org/licenses/LICENSE-2.0\n",
        "#\n",
        "# Unless required by applicable law or agreed to in writing, software\n",
        "# distributed under the License is distributed on an \"AS IS\" BASIS,\n",
        "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
        "# See the License for the specific language governing permissions and\n",
        "# limitations under the License."
      ]
    },
    {
      "cell_type": "markdown",
      "id": "e6aa9ed3-54c1-48e3-8c8f-e8f3cf152ec3",
      "metadata": {
        "id": "e6aa9ed3-54c1-48e3-8c8f-e8f3cf152ec3"
      },
      "source": [
        "# Using Gemini Function Calling to Get Bigquery Query Plan and Provide Optimization Recommendaiton using the Query Plan"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "7e82832d-fb80-4b04-9905-79bc6008d42c",
      "metadata": {
        "id": "7e82832d-fb80-4b04-9905-79bc6008d42c"
      },
      "source": [
        "<img width=\"80%\" src=\"\">"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "5a0b9005-e629-4206-b7ee-635d36423760",
      "metadata": {
        "id": "5a0b9005-e629-4206-b7ee-635d36423760"
      },
      "source": [
        "<img width=\"80%\" src=\"\">"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "56bad415-a400-42c3-800f-5c700a6c6cd7",
      "metadata": {
        "id": "56bad415-a400-42c3-800f-5c700a6c6cd7"
      },
      "source": [
        "## Getting Started\n",
        "\n",
        "### Install Vertex AI SDK and other required packages"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "ad3fc543-8dc7-43da-b095-0e4a965b7de4",
      "metadata": {
        "tags": [],
        "id": "ad3fc543-8dc7-43da-b095-0e4a965b7de4"
      },
      "outputs": [],
      "source": [
        "!pip install --upgrade google-cloud-aiplatform requests"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "5c810968-005f-4776-8d2b-99e04a49b550",
      "metadata": {
        "id": "5c810968-005f-4776-8d2b-99e04a49b550"
      },
      "source": [
        "### Restart runtime\n",
        "\n",
        "To use the newly installed packages in this Jupyter runtime, you must restart the runtime. You can do this by running the cell below, which restarts the current kernel.\n",
        "\n",
        "The restart might take a minute or longer. After it's restarted, continue to the next step."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "1e0edafe-5cf9-4979-87fa-79958402f9dc",
      "metadata": {
        "id": "1e0edafe-5cf9-4979-87fa-79958402f9dc"
      },
      "outputs": [],
      "source": [
        "import IPython\n",
        "\n",
        "app = IPython.Application.instance()\n",
        "app.kernel.do_shutdown(True)"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "299649d4-7959-455a-8971-d0d3bb5cc474",
      "metadata": {
        "id": "299649d4-7959-455a-8971-d0d3bb5cc474"
      },
      "source": [
        "### Authenticate your notebook environment (Colab only)\n",
        "\n",
        "If you are running this notebook on Google Colab, run the cell below to authenticate your environment."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "e641fb7d-2ed9-492e-ba8e-fabc86721630",
      "metadata": {
        "id": "e641fb7d-2ed9-492e-ba8e-fabc86721630"
      },
      "outputs": [],
      "source": [
        "import sys\n",
        "\n",
        "if \"google.colab\" in sys.modules:\n",
        "    from google.colab import auth\n",
        "\n",
        "    auth.authenticate_user(project_id=\"<project_id>\")"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "95c810db-eeea-4475-8c03-9016429272f8",
      "metadata": {
        "id": "95c810db-eeea-4475-8c03-9016429272f8"
      },
      "source": [
        "### Set Google Cloud project information and initialize Vertex AI SDK\n",
        "\n",
        "To get started using Vertex AI, you must have an existing Google Cloud project and [enable the Vertex AI API](https://console.cloud.google.com/flows/enableapi?apiid=aiplatform.googleapis.com).\n",
        "\n",
        "Learn more about [setting up a project and a development environment](https://cloud.google.com/vertex-ai/docs/start/cloud-environment)."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "529b3f2d-c270-4937-9b70-16651f260125",
      "metadata": {
        "id": "529b3f2d-c270-4937-9b70-16651f260125"
      },
      "outputs": [],
      "source": [
        "PROJECT_ID = \"<project_id>\"  # @param {type:\"string\"}\n",
        "LOCATION = \"us-central1\"  # @param {type:\"string\"}\n",
        "\n",
        "import vertexai\n",
        "\n",
        "vertexai.init(project=PROJECT_ID, location=LOCATION)"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "a0b04f91-62a4-48cf-8a84-50d41aaf5d79",
      "metadata": {
        "id": "a0b04f91-62a4-48cf-8a84-50d41aaf5d79"
      },
      "source": [
        "### Import libraries"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "51a1e99c-a72d-4f83-8c5c-f1e67592230d",
      "metadata": {
        "tags": [],
        "id": "51a1e99c-a72d-4f83-8c5c-f1e67592230d"
      },
      "outputs": [],
      "source": [
        "import requests\n",
        "from IPython.display import display, Markdown\n",
        "from vertexai.generative_models import (\n",
        "    Content,\n",
        "    FunctionDeclaration,\n",
        "    GenerativeModel,\n",
        "    GenerationConfig,\n",
        "    Part,\n",
        "    Tool,\n",
        ")\n",
        "from vertexai.preview.generative_models import ToolConfig"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Declare the function to retrieve the Bigquery Job Id"
      ],
      "metadata": {
        "id": "BmsRoW97qkZR"
      },
      "id": "BmsRoW97qkZR"
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "fd4dbed2-35f3-40b7-8bdf-abd598d948e7",
      "metadata": {
        "tags": [],
        "id": "fd4dbed2-35f3-40b7-8bdf-abd598d948e7"
      },
      "outputs": [],
      "source": [
        "get_optimization_results = FunctionDeclaration(\n",
        "    name=\"get_optimization_results\",\n",
        "    description=\"Provide specific and actionable optimization techniques to improve SQL query performance.\",\n",
        "    parameters={\n",
        "        \"type\": \"object\",\n",
        "        \"properties\": {\n",
        "            \"bq_job_id\":{\n",
        "                 \"type\": \"string\",\n",
        "                \"description\": \"The Bigquery Job id for which the optimization technique needs to be fetched\",\n",
        "            }\n",
        "        },\n",
        "    },\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "#Prompt for Getting the Optimization Output"
      ],
      "metadata": {
        "id": "rd5OY65oRN6J"
      },
      "id": "rd5OY65oRN6J"
    },
    {
      "cell_type": "code",
      "source": [
        "optimization_prompt = \"\"\"**Purpose and Goals:**\n",
        "* Analyze BigQuery execution statistics in JSON format to identify performance bottlenecks.\n",
        "* Provide specific and actionable optimization techniques to improve SQL query performance.\n",
        "**Behaviors and Rules:**\n",
        "1. **JSON Parsing:**\n",
        "* Load and parse the provided BigQuery execution statistics JSON.\n",
        "* Extract relevant metrics, including:\n",
        "* Total processing time\n",
        "* Slot usage\n",
        "* Bytes processed\n",
        "* Bytes shuffled\n",
        "* Compute vs. I/O wait time\n",
        "2. **Bottleneck Identification:**\n",
        "* Analyze the extracted metrics to identify potential performance bottlenecks, such as:\n",
        "* Large data scans (full table scans)\n",
        "* Inefficient join operations\n",
        "* Complex filtering or aggregation logic\n",
        "* Sub-optimal use of materialized views or cached results\n",
        "* Skewed data distribution\n",
        "3. **Optimization Recommendations:**\n",
        "* Based on the identified bottlenecks, provide tailored optimization techniques. These may include:\n",
        "* **Indexing:** Create appropriate indexes to reduce data scans.\n",
        "* **Partitioning and Clustering:** Organize large tables to improve query efficiency.\n",
        "* **Denormalization:** Consider strategic denormalization to reduce joins.\n",
        "* **Query Rewriting:** Restructure SQL queries for better performance (e.g., avoid subqueries, optimize JOIN types, leverage window functions).\n",
        "* **Data Filtering:** Filter data as early as possible in the query.\n",
        "* **Use of Materialized Views:** Pre-compute expensive aggregations or joins.\n",
        "* **Resource Adjustments:** Scale resources like slots or memory if needed.\n",
        "**Output Formatting:**\n",
        "* **Summary:** Present a clear summary of the analysis, including:\n",
        "* Key performance metrics\n",
        "* Identified bottlenecks\n",
        "* **Recommendations:** Provide a numbered list of actionable optimization recommendations, ordered by potential impact. Include brief explanations for each recommendation.\n",
        "**Overall Tone:**\n",
        "* Be informative and helpful.\n",
        "* Use clear and concise language.\n",
        "* Avoid assumptions about the user's technical expertise. Explain concepts as needed.\n",
        "\"\"\""
      ],
      "metadata": {
        "id": "D3I0QypLC2p3"
      },
      "id": "D3I0QypLC2p3",
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "id": "91bd0576",
      "metadata": {
        "id": "91bd0576"
      },
      "source": [
        "### Wrap function declarations in a tool\n",
        "\n",
        "Now, you can define a tool that will allow Gemini to select from the set of functions we've defined:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "6d5cd101-14d7-4562-8f5c-73a08fdc28b2",
      "metadata": {
        "tags": [],
        "id": "6d5cd101-14d7-4562-8f5c-73a08fdc28b2"
      },
      "outputs": [],
      "source": [
        "bigquery_insights_tool = Tool(\n",
        "    function_declarations=[\n",
        "        get_optimization_results\n",
        "    ],\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "bf483769-25ac-40f2-bf17-7da437f2e68a",
      "metadata": {
        "id": "bf483769-25ac-40f2-bf17-7da437f2e68a"
      },
      "source": [
        "#Generating API Key to Call the Bigquery Get Job API"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "import google.auth\n",
        "import google.auth.transport.requests\n",
        "import requests"
      ],
      "metadata": {
        "id": "0NRa0AN6Xagb"
      },
      "id": "0NRa0AN6Xagb",
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "creds, _ = google.auth.default()\n",
        "auth_req = google.auth.transport.requests.Request()\n",
        "creds.refresh(auth_req)"
      ],
      "metadata": {
        "id": "SquReOsGX4Jv"
      },
      "id": "SquReOsGX4Jv",
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "id": "6930cc77-577d-4dd5-998e-af0db2ca9971",
      "metadata": {
        "id": "6930cc77-577d-4dd5-998e-af0db2ca9971"
      },
      "source": [
        "### Define Python functions and a function handler"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "6a2ad181",
      "metadata": {
        "id": "6a2ad181"
      },
      "source": [
        "Python Function to call the Bigquery Get Job API"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "672ee652-a36f-4267-8e54-84bc3ad07bf7",
      "metadata": {
        "tags": [],
        "id": "672ee652-a36f-4267-8e54-84bc3ad07bf7"
      },
      "outputs": [],
      "source": [
        "def get_optimization_results(bq_job_id):\n",
        "    project_id=bq_job_id.split(\":\")[0]\n",
        "    location=bq_job_id.split(\":\")[1].split(\".\")[0]\n",
        "    job_id=bq_job_id.split(\":\")[1].split(\".\")[1]\n",
        "    url = f\"https://bigquery.googleapis.com/bigquery/v2/projects/\"+project_id+\"/jobs/\"+job_id\n",
        "    api_request = requests.get(url,headers={'Authorization': f'Bearer {creds.token}'},params={\"location\":location})\n",
        "    print(api_request.json())\n",
        "    return api_request.json()"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "16bde708-dfbb-46e9-8070-d8ef687c9a06",
      "metadata": {
        "id": "16bde708-dfbb-46e9-8070-d8ef687c9a06"
      },
      "source": [
        "Define a function handler that maps function call names (from your function declarations) to actual Python functions that call APIs:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "c2594766-4f93-4dbf-982b-c50131a9fbd9",
      "metadata": {
        "id": "c2594766-4f93-4dbf-982b-c50131a9fbd9"
      },
      "outputs": [],
      "source": [
        "function_handler = {\n",
        "    \"get_optimization_results\": get_optimization_results\n",
        "}"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "tool_config = ToolConfig(\n",
        "    function_calling_config=ToolConfig.FunctionCallingConfig(\n",
        "        mode=ToolConfig.FunctionCallingConfig.Mode.ANY\n",
        "    )\n",
        ")"
      ],
      "metadata": {
        "id": "py35GdGD7FtM"
      },
      "id": "py35GdGD7FtM",
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "id": "25bf5867",
      "metadata": {
        "id": "25bf5867"
      },
      "source": [
        "### Initialize model\n",
        "\n",
        "Initialize the Gemini model with the desired model parameters and `Tool` that we defined earlier:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "9938f4ec-edd7-4738-85bd-ed125bdd54c4",
      "metadata": {
        "tags": [],
        "id": "9938f4ec-edd7-4738-85bd-ed125bdd54c4"
      },
      "outputs": [],
      "source": [
        "gemini_model = GenerativeModel(\n",
        "    \"gemini-1.5-pro-001\",\n",
        "    generation_config=GenerationConfig(temperature=0),\n",
        "    tools=[bigquery_insights_tool]\n",
        ")"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "chat = gemini_model.start_chat()\n"
      ],
      "metadata": {
        "id": "RJQgJD7bBLRf"
      },
      "id": "RJQgJD7bBLRf",
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "#Defining the function for sending the message and getting the response back"
      ],
      "metadata": {
        "id": "L7OzwIHwRrO-"
      },
      "id": "L7OzwIHwRrO-"
    },
    {
      "cell_type": "code",
      "source": [
        "def send_chat_message(prompt):\n",
        "    display(Markdown(\"#### Prompt\"))\n",
        "    print(prompt, \"\\n\")\n",
        "    # Send a chat message to the Gemini API\n",
        "    if not isinstance(prompt, str):\n",
        "      raise TypeError(\"The 'prompt' variable must be a string.\")\n",
        "    if not prompt:\n",
        "      raise ValueError(\"The 'prompt' variable cannot be empty.\")\n",
        "\n",
        "    response = chat.send_message(prompt+optimization_prompt)\n",
        "\n",
        "    # Extract the function call response\n",
        "    function_call = response.candidates[0].content.parts[0].function_call\n",
        "\n",
        "    # Extract the function call name\n",
        "    function_name = function_call.name\n",
        "    display(Markdown(\"#### Predicted function name\"))\n",
        "    print(function_name, \"\\n\")\n",
        "\n",
        "    # Extract the function call parameters\n",
        "    params = {key: value for key, value in function_call.args.items()}\n",
        "    display(Markdown(\"#### Predicted function parameters\"))\n",
        "    print(params, \"\\n\")\n",
        "\n",
        "    # Invoke a function that calls an external API\n",
        "    function_api_response = function_handler[function_name](params['bq_job_id'])\n",
        "    display(Markdown(\"#### API response\"))\n",
        "    print(function_api_response)\n",
        "\n",
        "    response = chat.send_message(\n",
        "                Part.from_function_response(\n",
        "                    name=function_name,\n",
        "                    response={\"content\": function_api_response},\n",
        "                ),\n",
        "            )\n",
        "\n",
        "    # Show the final natural language summary\n",
        "    display(Markdown(\"#### Natural language response\"))\n",
        "    display(Markdown(response.text.replace(\"$\", \"\\\\\\$\")))"
      ],
      "metadata": {
        "id": "l9s5lGHLYiSj"
      },
      "id": "l9s5lGHLYiSj",
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [],
      "metadata": {
        "id": "ynowuXBPR4fz"
      },
      "id": "ynowuXBPR4fz"
    },
    {
      "cell_type": "markdown",
      "source": [
        "#Sample Query from User by passing the Bigquery Job Id"
      ],
      "metadata": {
        "id": "4APkFecKR55y"
      },
      "id": "4APkFecKR55y"
    },
    {
      "cell_type": "code",
      "source": [
        "send_chat_message(\"How to optimize the bigquery job id <project_id>:US.bquxjob_5712d3a6_18fb981a2b0\")"
      ],
      "metadata": {
        "id": "Dtr9U56rD1F8"
      },
      "id": "Dtr9U56rD1F8",
      "execution_count": null,
      "outputs": []
    }
  ],
  "metadata": {
    "environment": {
      "kernel": "conda-root-py",
      "name": "workbench-notebooks.m115",
      "type": "gcloud",
      "uri": "gcr.io/deeplearning-platform-release/workbench-notebooks:m115"
    },
    "kernelspec": {
      "display_name": "Python 3 (ipykernel)",
      "language": "python",
      "name": "python3"
    },
    "language_info": {
      "codemirror_mode": {
        "name": "ipython",
        "version": 3
      },
      "file_extension": ".py",
      "mimetype": "text/x-python",
      "name": "python",
      "nbconvert_exporter": "python",
      "pygments_lexer": "ipython3",
      "version": "3.11.8"
    },
    "colab": {
      "provenance": []
    }
  },
  "nbformat": 4,
  "nbformat_minor": 5
}